import numpy as np 
import matplotlib.pyplot as plt 
from os import path

import config

data_folder = 'output'
dataset_name = config.DATASET_NAME
sample_num = config.TEST_NUM
graph_rate = config.LOG_RATE
btch_size = config.BATCH_SIZE

def plot_data(mode='loss_data', names=None, signs=None, stride=1):
    ''' plot the data from npy file
    '''
    if signs is None:
        signs = ['-'] * len(names)
    elif len(names) != len(signs):
        print('(!) The length of two iterable object do not match.')
        exit()
    if mode not in ('loss_data', 'acc_data', 'epoch_acc'):
        print('Unsupported data mode: %s' % mode)
        exit()
    plt.figure()
    legends = []
    for name, sign in zip(names, signs):
        data = np.load(path.join(
            data_folder,
            mode + '-' + name.lower() + '.npy'), allow_pickle=True)
        if mode == 'epoch_acc':
            plt.plot(data[::stride], sign)
        else:
            plt.plot(data[0][::stride], data[1][::stride], sign)
        legends.append(name)
    if mode == 'loss_data':
        plt.title('Training loss on %s' % dataset_name)
        plt.xlabel('training progress (samples)')
        plt.ylabel('training accuracy (%)')
    elif mode == 'acc_data':
        plt.title('Training accuracy on %s' % dataset_name)
        plt.xlabel('training progress (samples)')
        plt.ylabel('accuracy on training set (%)')
    else:
        plt.title('Validating accuracy on %s' % dataset_name)
        plt.xlabel('training progress (epoch num.)')
        plt.ylabel('accuracy on validation set (%)')
    plt.legend(legends)
    # plt.show()

def compare_acc(name, stride=1):
    ''' draw accuracy on train and val set in the same figure
    '''
    plt.figure()
    train_acc = np.load(path.join(
        data_folder,
        'acc_data-' + name.lower() + '.npy'
    ), allow_pickle=True)
    val_acc = np.load(path.join(
        data_folder,
        'epoch_acc-' + name.lower() + '.npy'
    ), allow_pickle=True)
    val_count = np.array(range(len(val_acc))) * config.TRAIN_NUM
    plt.plot(train_acc[0][::stride], train_acc[1][::stride])
    plt.plot(val_count, val_acc)
    plt.title("%s's accuracy on %s" % (name, dataset_name))
    plt.xlabel('training progress (samples)')
    plt.ylabel('accuracy (%)')
    plt.legend(['train', 'val'])



if __name__ == '__main__':
    names = ['VGG16']
    signs = ['-+']
    # plot_data('loss_data', names, signs)
    # plot_data('acc_data', names, signs)
    # plot_data('epoch_acc', names, signs)
    compare_acc('VGG16', stride=5)
    plt.show()
    
    